from lagent.prompts import StrParser
from lagent.agents.aggregator import DefaultAggregator
from typing import List, Union
from lagent.memory import Memory

class FewshotAggregator(DefaultAggregator):
    """支持few-shot示例的消息聚合器
    
    继承自DefaultAggregator,添加了对few-shot示例的支持
    """
    def __init__(self, few_shot: List[dict] = None):
        """
        Args:
            few_shot: few-shot示例列表,每个示例是包含role和content的字典
        """
        self.few_shot = few_shot or []

    def aggregate(self,
                  messages: Memory,
                  name: str, 
                  parser: StrParser = None,
                  system_instruction: Union[str, dict, List[dict]] = None) -> List[dict]:
        """聚合消息历史、系统提示和few-shot示例
        
        Args:
            messages: 对话历史记忆
            name: 智能体名称
            parser: 消息解析器
            system_instruction: 系统提示
            
        Returns:
            聚合后的消息列表,符合OpenAI格式
        """
        _message = []
        
        # 添加系统提示
        if system_instruction:
            _message.extend(
                self.aggregate_system_intruction(system_instruction))
                
        # 添加few-shot示例
        _message.extend(self.few_shot)
        
        # 获取历史消息
        messages = messages.get_memory()
        
        # 处理每条历史消息
        for message in messages:
            if message.sender == name:
                # 智能体发送的消息作为assistant角色
                _message.append(
                    dict(role='assistant', content=str(message.content)))
            else:
                # 用户发送的消息作为user角色
                user_message = message.content
                # 如果上一条也是用户消息,则合并内容
                if len(_message) > 0 and _message[-1]['role'] == 'user':
                    _message[-1]['content'] += user_message
                else:
                    _message.append(dict(role='user', content=user_message))
        return _message
